// SPDX-FileCopyrightText: 2023-2025 Sangfor Technologies Inc.
// SPDX-License-Identifier: Mulan PSL v2
package com.geniusai.aip.auth.config;

import com.geniusai.aip.auth.domain.ResourceGpu;
import com.geniusai.aip.auth.domain.ResourceGpuDO;
import com.geniusai.aip.auth.mapper.ResourceGpuMapper;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;

import javax.annotation.PostConstruct;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author lingchuanzhou19320
 * @date 2022/6/6
 */
@Configuration
public class GpuResourceConfig {

    @Autowired
    ResourceGpuMapper resourceGpuMapper;

    private static Map<String, ResourceGpu> gpuResources = new HashMap<>();

    @PostConstruct
    public void initGpuResourceConfig() {
        List<ResourceGpuDO> resourceGpuDOS = resourceGpuMapper.selectList(null);
        for (ResourceGpuDO resourceGpuDO : resourceGpuDOS) {
            ResourceGpu resourceGpu = new ResourceGpu();
            BeanUtils.copyProperties(resourceGpuDO, resourceGpu);
            gpuResources.put(resourceGpu.getGpuType(), resourceGpu);
        }
    }

    public static List<ResourceGpu> getGpuResources() {
        return gpuResources.values().stream().collect(Collectors.toList());
    }

    public static Integer getTotalGpuMemory() {
        Integer totalGpuMemory = 0;
        for (String gpuType: gpuResources.keySet()) {
            totalGpuMemory += gpuResources.get(gpuType).getGpuSize() * gpuResources.get(gpuType).getGpuNum();
        }
        return totalGpuMemory;
    }

    public static ResourceGpu getGpuResourceByGpuType(String gpuType) {
        if (StringUtils.isEmpty(gpuType)) {
            ResourceGpu resourceGpu = new ResourceGpu();
            resourceGpu.setGpuNum(0);
            resourceGpu.setGpuSize(0);
            resourceGpu.setGpuType("");
            resourceGpu.setId(0);
            resourceGpu.setVGpuNum(0);
            resourceGpu.setVGpuRation(0);
            resourceGpu.setVGpuSize(0);
            return resourceGpu;
        }
        return gpuResources.get(gpuType);
    }

    public static Integer getVGpuRatioByGpuType(String gpuType) {
        if (StringUtils.isEmpty(gpuType)) {
            return 0;
        }
        return gpuResources.get(gpuType).getVGpuRation();
    }
}
